
package com.microsoft.mcp.sample.server.service;

import com.fasterxml.jackson.annotation.JsonProperty;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.context.annotation.Description;
import org.springframework.web.reactive.function.client.WebClient;
import reactor.core.publisher.Mono;
import org.springframework.beans.factory.annotation.Value;

import java.util.List;
import java.util.Map;
import java.util.function.Function;

/**
 * Configuration class that defines the MCP Knowledge Base tools as Spring Beans.
 * Each bean is a function that the Spring AI agent can discover and execute.
 * The @Description annotation is critical for the agent to understand the purpose of each tool.
 */
@Configuration
public class McpKnowledgeTools {

    @Value("${mcp.server.base-url}")
    private String mcpBaseUrl;

    /**
     * Creates a WebClient bean configured to communicate with the MCP server.
     *
     * @return A configured WebClient instance.
     */
    @Bean
    public WebClient mcpWebClient() {
        return WebClient.builder()
                .baseUrl(mcpBaseUrl) // MCP server address
                .defaultHeader("Content-Type", "application/json")
                .build();
    }

    @Bean
    @Description("Search for content in the knowledge base that is semantically similar to the user's query. Use this for questions, concepts, or finding related information.")
    public Function<SearchSimilarRequest, Mono<McpSearchResponse>> searchSimilarInKnowledgeBase(WebClient mcpWebClient) {
        return request -> {
            System.out.println("Agent is calling searchSimilarInKnowledgeBase with query: " + request.query);
            return mcpWebClient.post()
                    .uri("/api/tools/search_similar")
                    .bodyValue(request)
                    .retrieve()
                    .bodyToMono(McpSearchResponse.class);
        };
    }

    @Bean
    @Description("Add or update a text document in the knowledge base. Use this when the user wants to save, store, or remember a piece of text.")
    public Function<VectorizeTextRequest, Mono<McpVectorizeResponse>> addTextToKnowledgeBase(WebClient mcpWebClient) {
        return request -> {
            System.out.println("Agent is calling addTextToKnowledgeBase for source_path: " + request.sourcePath);
            // In a real application, source_path should be unique and meaningful.
            var enrichedRequest = new VectorizeTextRequest(
                request.content,
                request.sourcePath != null ? request.sourcePath : "agent-doc-" + System.currentTimeMillis(),
                "spring_ai_agent" // Use a dedicated table for the agent
            );
            return mcpWebClient.post()
                    .uri("/api/tools/vectorize_text")
                    .bodyValue(enrichedRequest)
                    .retrieve()
                    .bodyToMono(McpVectorizeResponse.class);
        };
    }
    
    // Add other MCP tools like vectorize_code, list_indexed_content etc. here following the same pattern.
    // For brevity, only search_similar and vectorize_text are implemented in this example.


    // --- DTO (Data Transfer Object) Records for MCP API ---
    // These records define the request and response structures for our tool functions.
    // The JSON properties match the MCP server's tool parameters.

    public record SearchSimilarRequest(
        @JsonProperty(required = true) String query,
        @JsonProperty("top_k") Integer topK,
        @JsonProperty("table_name") String tableName
    ) {
        public SearchSimilarRequest(String query) {
            this(query, 5, "spring_ai_agent");
        }
    }

    public record VectorizeTextRequest(
        @JsonProperty(required = true) String content,
        @JsonProperty("source_path") String sourcePath,
        @JsonProperty("table_name") String tableName
    ) {
        public VectorizeTextRequest(String content, String sourcePath) {
            this(content, sourcePath, "spring_ai_agent");
        }
    }

    public record McpVectorizeResponse(boolean success, String message, @JsonProperty("chunks_created") int chunksCreated) {}
    public record McpSearchResponse(boolean success, String message, List<McpSearchResult> results) {}
    public record McpSearchResult(String content, double score, @JsonProperty("source_path") String sourcePath) {}
} 